#!/usr/bin/python2
# -*- coding: utf-8 -*-
# Name:     ms.py
# Author:   xiooli <xioooli[at]yahoo.com.cn>
# Version:  111109

import os, re, types, json, time
from pattern import forge_pattern
from mzxml import MzXML
from .. mol import Mol, exact_mass, mol_weight
from .. import generator

regx_space = re.compile(r'[ \t]+')

def svg_line(x1, y1, x2, y2):
    return '<line x1="%s" y1="%s" x2="%s" y2="%s"/>\n' %(x1, y1, x2, y2)

def parse_mzXML(file_name):
    assert os.access(file_name, os.R_OK)
    mzXML = MzXML()
    mzXML.parse_file(file_name)
    return mzXML.MS_list

def extract_MSn(ms_list, ms_level):
    rst_ms_list = []
    for ms in ms_list:
        if ms.ms_level == ms_level:
            rst_ms_list.append(ms)
    return rst_ms_list

class MS():
    def __init__(self, precursor = None, adduct = 1, spectra = ''):
        if os.path.exists(spectra):
            self.spectra_list = [ l.strip() for l in open(spectra, 'r').readlines() ]
        else:
            self.spectra_list = spectra.replace('\r', '').split('\n')

        is_valid_sp = False

        self.precursor = precursor
        self.adduct = adduct
        self.repl_dict = {
                # ME, mass of electron
                'ME': 0.000549,
                'MA': self.precursor,
                'AI': self.adduct,
                'M': 0,
                }

        for line in self.spectra_list:
            if line.startswith('Mass'):
                is_valid_sp = True
                break
        if not is_valid_sp:
            self.spectra_list.insert(0, 'Mass')

        self.spectra_txt = '\n'.join(self.spectra_list)
        self.noise_level = 0.05 # 5 percent
        self.iso_width = 0.5
        self.round = 4
        self.mz_list = []
        self.mz_intenpct_list = []
        self.neutral_losses = []
        self.resolution = 0 # 0 for low resolution MS, 1 for high resolution MS
        self.deviation = 5 / 1000000.0 # 5 ppm

    def set_precursor(self, precursor):
        self.precursor = precursor
        self.repl_dict['MA'] = self.precursor
        self.repl_dict['M'] = self.precursor - self.adduct

    def set_adduct(self, adduct):
        self.adduct = adduct
        self.repl_dict['AI'] = self.adduct
        self.repl_dict['M'] = self.precursor - self.adduct

    def get_ms_list(self):
        '''Parses the (mz, intensity) list'''

        if self.mz_intenpct_list:
            return

        rec = False
        mz_list = []
        int_list = []

        for line in self.spectra_list:
            line = regx_space.sub('\t', line.strip('\r\n\t'))
            if not line:
                continue
            if not rec and line.startswith('Mass'):
                rec = True
            elif rec:
                try:
                    mz, inten = line.split('\t')
                except:
                    pass
                mz = round(float(mz), self.round)
                mz_list.append(mz)
                int_list.append(round(float(inten),2))

        max_inten = max(int_list)
        pct_list = [ round(100*i/max_inten, 2) for i in int_list ]

        for i in zip(mz_list, pct_list):
            if i[1] >= self.noise_level * 100:
                self.mz_list.append(i[0])
                self.mz_intenpct_list.append(i)
        return self.mz_intenpct_list

    def gen_comb(self, n, peaks, sorted = False):
        '''Generats all the combinations of n members out of peaks'''
        if not sorted:
            peaks.sort()

        if n > 0:
            for i in range(len(peaks)):
                v = peaks[i:i+1]
                if n == 1:
                    yield v
                else:
                    rest = peaks[i+1:]
                    for c in self.gen_comb(n - 1, rest, True):
                        yield v + c
        else:
            yield []

    def get_neutral_losses(self):
        '''Get the neutral loss list'''
        if self.neutral_losses:
            return self.neutral_losses

        if not self.mz_list:
            self.get_ms_list()

        peaks = self.mz_list[::-1]

        for p1, p2 in self.gen_comb(2, peaks):
            self.neutral_losses.append(p1-p2)

        return self.neutral_losses

    def exist_neutral_losses(self, *losses):
        '''Check if the neutral losses exists'''
        if not self.neutral_losses:
            self.get_neutral_losses()

        results = []

        for loss in losses:
            for nl in self.neutral_losses:
                if self.resolution == 0:
                    if abs(loss - nl)  <= self.iso_width:
                        results.append(nl)
                        break
                else:
                    if abs((loss - nl) / float(loss)) <= self.deviation:
                        results.append(nl)
                        break

        if len(results) == len(losses):
            return True
        return False

    def exist_nl(self, *losses):
        '''Alisa of method exist_neutral_losses'''
        return self.exist_neutral_losses(*losses)

    def exist_nls(self, *losses):
        '''Alisa of method exist_neutral_losses'''
        return self.exist_neutral_losses(*losses)

    def equal(self, p1, p2):
        '''Check if p1 equals p2'''
        if self.resolution == 0:
            #print p1, p2, p1-p2
            if abs(p1 - p2) <= sel/f.iso_width:
                return True
        elif self.resolution == 1:
            if abs((p1-p2) / float(p1)) <= self.deviation:
                return True
        return False

    def exist_peaks(self, *peaks):
        '''Check if the given peaks exist. Returns a list contains the acurate
        mz and intensity of a given peak in a tuple.'''

        if not self.mz_intenpct_list:
            self.get_ms_list()

        results = []

        for peak in peaks:
            for mz_inten in self.mz_intenpct_list:
                if self.resolution == 0:
                    if abs(peak - mz_inten[0]) <= self.iso_width:
                        results.append(mz_inten)
                        break
                else:
                    if abs((peak - mz_inten[0]) / float(peak)) <= self.deviation:
                        results.append(mz_inten)
                        break
        if len(results) == len(peaks):
            return True
        return False

    def exist(self, *peaks):
        '''Alias of method exist'''
        return self.exist_peaks(*peaks)

    def exists(self, *peaks):
        '''Alias of method exist'''
        return self.exist_peaks(*peaks)

    def intensity(self, peak):
        '''Returns the intensity of the given peak if exists'''

        if not self.mz_intenpct_list:
            self.get_ms_list()

        for mz_inten in self.mz_intenpct_list:
            if self.resolution == 0:
                if abs(peak - mz_inten[0]) <= self.iso_width:
                    return mz_inten[1]
            else:
                if abs((peak - mz_inten[0]) / float(peak)) <= self.deviation:
                    return mz_inten[1]
        return 0

    def inten(self, peak):
        '''alisa of method intensity'''
        return self.intensity(peak)

    def analysis(self, pattern):
        '''Apply the given pattern and return the results of the
        application of the pattern to this MS'''

        # results should named the same as in the pattern_code
        t0 = time.time()
        if not (pattern) or (not self.precursor) or (not self.adduct):
            return []

        if type(pattern) == types.StringType:
            pattern = json.loads(pattern)

        if not self.mz_list:
            self.get_ms_list()

        peaks = self.mz_list[::-1]
        len_args, pattern_code = forge_pattern(pattern)

        pattern_func = compile(pattern_code, '<string>', 'exec')

        results = []
        args_dict = {
                    'self': self,
                    'MS': self,
                    'Mol': Mol,
                    'Generator': generator,
                    'G': generator,
                    'Exact_mass': exact_mass,
                    'EM': exact_mass,
                    'Mol_weight': mol_weight,
                    'MW': mol_weight,
                    'Results': results,
                    'R': results,
                }
        args_dict.update(self.repl_dict)

        for i in self.gen_comb(len_args, peaks):
            args_dict['args'] = i
            eval(pattern_func, args_dict)

        #print results
        tmpl = []
        for l in results:
            if not l in tmpl:
                tmpl.append(l)
        return tmpl

    def gen_ms_svg(self, width = 1000, height = 500):
        '''Genterats the svg presents this ms'''
        if not self.mz_intenpct_list:
            self.get_ms_list()

        max_mz = max(self.mz_list)
        max_mz = max_mz/100 > round(max_mz/100) and (round(max_mz/100) + 1)*100 \
                or round(max_mz/100)*100
        step = 100
        while max_mz/step > 10:
            step += 50

        svg = '''<?xml version="1.0" encoding="UTF-8"?>
<svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" width="%s" height="%s" viewBox="0 -30 1000 500" version="1.1">
<defs>
    <style type="text/css"><![CDATA[
        line {
            fill: black;
            stroke: black;
            stroke-width: 3;
            }
        text {
            fill: black;
            font-family: Verdana;
            font-size: 30px;
            font-weight: normal;
            }]]>
    </style>
</defs>\n''' %(width, height)
        x0 = 85
        x1 = 80
        y0 = height - x1

        svg += '<text x="40" y="25">RA</text>\n'

        # y label
        svg += svg_line(x0, y0, x0, x1)
        for i in range(6):
            y = y0 - (height-x1*2)/5.0 * i
            svg += svg_line(x0 - 15, y, x0+1, y)
            svg += '<text x="10" y="%s">%s</text>\n' %(y+8, i*20)

        # x label
        svg += svg_line(x0, y0, width - x0, y0)
        svg += '<text x="%s" y="%s">m/z</text>\n' %(width-x0 + 10, y0 + 5)
        for i in range(int(max_mz/step) + 1):
            x = x0 + (width-2*x0)/float(max_mz/step) * i
            svg += svg_line(x, y0 - 1, x, y0 + 15)
            if i == 0:
                svg += '<text x="%s" y="%s">%s</text>\n' %(x - 10, y0 + 40, i*step)
            else:
                svg += '<text x="%s" y="%s">%s</text>\n' %(x - 28, y0 + 40, i*step)

        # peak with m/z label
        for mz, inten in self.mz_intenpct_list:
            if self.resolution == 0:
                dis_mz = int(round(mz, 0))
            elif self.resolution == 1:
                dis_mz = round(mz, 4)

            x = x0 + mz/step*float(width-2*x0)/float(max_mz/step)
            y = y0 - (inten/100.0)*(height-x1*2)
            svg += svg_line(x, y0, x, y)
            svg += '<g transform="translate(%s, %s) rotate(-90)"><text style="font-size:27">%s</text></g>\n' \
                %(x + 8, y - 3, dis_mz)

        svg += '</svg>'

        return svg

if __name__ == '__main__':

    pattern = open('tst.txt').read()

    ms = MS(spectra = 'ms.txt')
    ms.resolution = 1
    ms.set_precursor(655)
    ms.set_adduct(22.98977)
    print ms.exist(655.1201-18)
    #print ms.get_ms_list()
    #print 'Exist peaks:', ms.exist(637.1234, 219, 323, 341, 405, 423, 459, 515, 533, 655)
    #print 'Intensity of 637.1234:', ms.inten(637.1234)
    #print 'Neutral losses:', ms.get_neutral_losses()
    #print 'Exist neutral loss 418?', ms.exist_nls(418)
    #print ms.gen_ms_svg()
    m=Mol()
    m.add_atom('*')
    for i in range(5):
        m.add_atom('C')
        m.add_bond(i, -1)
    s=m.get_mol()
    print s

